

#include "bezier_opt/trajectory_generator.h"

using namespace std;    
using namespace Eigen;

//static void MSKAPI printstr(void *handle, MSKCONST char str[])
//{
//  printf("%s",str);
//}

int TrajectoryGenerator::BezierPloyCoeffGeneration(
            const vector<Cube> &corridor,
            const MatrixXd &MQM,
            const MatrixXd &pos,
            const MatrixXd &vel,
            const MatrixXd &acc,
            const double maxVel,
            const double maxAcc,
            const int traj_order,
            const double minimize_order,
            const double margin,
            const bool & isLimitVel,
            const bool & isLimitAcc,
            double & obj,
            MatrixXd & PolyCoeff)  // define the order to which we minimize.   1 -- velocity, 2 -- acceleration, 3 -- jerk, 4 -- snap  
{   
#define ENFORCE_VEL  isLimitVel // whether or not adding extra constraints for ensuring the velocity feasibility
#define ENFORCE_ACC  isLimitAcc // whether or not adding extra constraints for ensuring the acceleration feasibility

    double initScale = corridor.front().t;
    double lstScale  = corridor.back().t;
    int segment_num  = corridor.size();

    int n_poly = traj_order + 1;
    int& s1d1CtrlP_num = n_poly;
    int s1CtrlP_num   = 3 * s1d1CtrlP_num;

    int equ_con_s_num = 3 * 3; // p, v, a in x, y, z axis at the start point
    int equ_con_e_num = 3 * 3; // p, v, a in x, y, z axis at the end point
    int equ_con_continuity_num = 3 * 3 * (segment_num - 1);
    int equ_con_num   = equ_con_s_num + equ_con_e_num + equ_con_continuity_num; // p, v, a in x, y, z axis in each segment's joint position
    
    int vel_con_num = 3 *  traj_order * segment_num;
    int acc_con_num = 3 * (traj_order - 1) * segment_num;

    if( !ENFORCE_VEL )
        vel_con_num = 0;

    if( !ENFORCE_ACC )
        acc_con_num = 0;

    int high_order_con_num = vel_con_num + acc_con_num; 
//    int high_order_con_num = 0; //3 * traj_order * segment_num;

    int con_num   = equ_con_num + high_order_con_num;
    int ctrlP_num = segment_num * s1CtrlP_num;

//    double x_var[ctrlP_num];
//    double primalobj;

    USING_NAMESPACE_QPOASES;
    int_t& nv = ctrlP_num;
    int_t& nc = con_num;
    unique_ptr<real_t[]> H(new real_t[nv*nv]);
    unique_ptr<real_t[]> A(new real_t[nv*nc]);
    unique_ptr<real_t[]> g(new real_t[nv]);
    unique_ptr<real_t[]> lb(new real_t[nv]);
    unique_ptr<real_t[]> ub(new real_t[nv]);
    unique_ptr<real_t[]> lbA(new real_t[nc]);
    unique_ptr<real_t[]> ubA(new real_t[nc]);

    for(int i = 0; i < nv*nv; i++)
        H[i] = 0;
    for(int i = 0; i < nv*nc; i++)
        A[i] = 0;
    for(int i = 0; i < nv; i++){
        g[i] = 0;
        lb[i] = 0;
        ub[i] = 0;
    }
    for(int i = 0; i < nc; i++){
        lbA[i] = 0;
        ubA[i] = 0;
    }

    int idx_c = 0;
    if(ENFORCE_VEL)
    {
        /***  Stack the bounding value for the linear inequality for the velocity constraints  ***/
        for(int i = 0; i < vel_con_num; i++){
            lbA[idx_c] = -maxVel;
            ubA[idx_c] = maxVel;
            idx_c++;
        }
    }

    if(ENFORCE_ACC)
    {
        /***  Stack the bounding value for the linear inequality for the acceleration constraints  ***/
        for(int i = 0; i < acc_con_num; i++){
            lbA[idx_c] = -maxAcc;
            ubA[idx_c] = maxAcc;
            idx_c++;
        }
    }

    for(int i = 0; i < equ_con_num; i ++ ){
        double beq_i;
        if(i < 3)                    beq_i = pos(0, i);
        else if (i >= 3  && i < 6  ) beq_i = vel(0, i - 3);
        else if (i >= 6  && i < 9  ) beq_i = acc(0, i - 6);
        else if (i >= 9  && i < 12 ) beq_i = pos(1, i - 9 );
        else if (i >= 12 && i < 15 ) beq_i = vel(1, i - 12);
        else if (i >= 15 && i < 18 ) beq_i = acc(1, i - 15);
        else beq_i = 0.0;

        // # cb_eq means: constriants boundary of equality constrain
        lbA[idx_c] = beq_i;
        ubA[idx_c] = beq_i;
        idx_c++;
//        std::cout << "put con:" << i << "," << beq_i << std::endl;
    }

    /* ## define a container for control points' boundary and boundkey ## */ 
    /* ## dataType in one tuple is : boundary type, lower bound, upper bound ## */

    int idx = 0;
    for(int k = 0; k < segment_num; k++)
    {   
        Cube cube_     = corridor[k];
        double scale_k = cube_.t;

        for(int i = 0; i < 3; i++ )
        {
//            std::cout << k << "," << i << " pos con : "
//            << (cube_.box[i].first  + margin) << " , "
//            << (cube_.box[i].second - margin) << std::endl;
            for(int j = 0; j < n_poly; j ++ )
            {
                double lo_bound, up_bound;
                if(k > 0)
                {
                    // i == 2 对应轴方向
                    if(i != 2) {
                        lo_bound = (cube_.box[i].first  + margin) / scale_k;
                        up_bound = (cube_.box[i].second - margin) / scale_k;
                    } else {
                        lo_bound = (cube_.box[i].first)  / scale_k;
                        up_bound = (cube_.box[i].second) / scale_k;
                    }

                }
                else
                {
                    lo_bound = (cube_.box[i].first)  / scale_k;
                    up_bound = (cube_.box[i].second) / scale_k;
                }
//                std::cout << "[traj]: seg, " << k << ", " << i << ", " << j << std::endl;
//                std::cout << "pos con : " << lo_bound * scale_k << ", " << up_bound * scale_k << std::endl;
                // # vb_x means: varialbles boundary of unknowns x (Polynomial coeff)
                lb[idx] = lo_bound;
                ub[idx] = up_bound;
                idx++;
            }
        } 
    }

// Parameters used in the optimizer
    int row_idx = 0;
    // The velocity constraints
    if(ENFORCE_VEL)
    {   
        for(int k = 0; k < segment_num ; k ++ )
        {   
            for(int i = 0; i < 3; i++)
            {  // for x, y, z loop
                for(int p = 0; p < traj_order; p++)
                {
                    //number of non-zero index
                    int nzi = 2;
                    unsigned int asub[nzi];
                    double aval[nzi];

                    aval[0] = -1.0 * traj_order;
                    aval[1] =  1.0 * traj_order;

                    asub[0] = k * s1CtrlP_num + i * s1d1CtrlP_num + p;    
                    asub[1] = k * s1CtrlP_num + i * s1d1CtrlP_num + p + 1;    

                    for(int i = 0; i < nzi; i++)
                        A[row_idx*nv + asub[i]] = aval[i];
                    row_idx++;
                }
            }
        }
    }

    // The acceleration constraints
    if(ENFORCE_ACC)
    {
        for(int k = 0; k < segment_num ; k ++ )
        {
            for(int i = 0; i < 3; i++)
            { 
                for(int p = 0; p < traj_order - 1; p++)
                {    
                    int nzi = 3;
                    unsigned int asub[nzi];
                    double aval[nzi];

                    aval[0] =  1.0 * traj_order * (traj_order - 1) / corridor[k].t;
                    aval[1] = -2.0 * traj_order * (traj_order - 1) / corridor[k].t;
                    aval[2] =  1.0 * traj_order * (traj_order - 1) / corridor[k].t;
                    asub[0] = k * s1CtrlP_num + i * s1d1CtrlP_num + p;    
                    asub[1] = k * s1CtrlP_num + i * s1d1CtrlP_num + p + 1;    
                    asub[2] = k * s1CtrlP_num + i * s1d1CtrlP_num + p + 2;

                    for(int i = 0; i < nzi; i++)
                        A[row_idx*nv + asub[i]] = aval[i];
                    row_idx++;
                }
            }
        }
    }

//    cout << ctrlP_num << "," << (high_order_con_num+ctrlP_num)*2 << "\n";
//    cout << "size of CI :" << CI.ncols() << "," << CI.nrows() << "\nsize of ci0 :" << ci0.size() << endl;
    /*   Start position  */
//    row_idx = 0;
    {
        // position :
        for(int i = 0; i < 3; i++)
        {  // loop for x, y, z       
            int nzi = 1;
            unsigned int asub[nzi];
            double aval[nzi];
            aval[0] = 1.0 * initScale;
            asub[0] = i * s1d1CtrlP_num;

            A[row_idx*nv + asub[0]] = aval[0];
            row_idx ++;
        }
        // velocity :
        for(int i = 0; i < 3; i++)
        {  // loop for x, y, z       
            int nzi = 2;
            unsigned int asub[nzi];
            double aval[nzi];
            aval[0] = - 1.0 * traj_order;
            aval[1] =   1.0 * traj_order;
            asub[0] = i * s1d1CtrlP_num;
            asub[1] = i * s1d1CtrlP_num + 1;
            for(int i = 0; i < nzi; i++)
                A[row_idx*nv + asub[i]] = aval[i];
            row_idx++;
        }
        // acceleration : 
        for(int i = 0; i < 3; i++)
        {  // loop for x, y, z       
            int nzi = 3;
            unsigned int asub[nzi];
            double aval[nzi];
            aval[0] =   1.0 * traj_order * (traj_order - 1) / initScale;
            aval[1] = - 2.0 * traj_order * (traj_order - 1) / initScale;
            aval[2] =   1.0 * traj_order * (traj_order - 1) / initScale;
            asub[0] = i * s1d1CtrlP_num;
            asub[1] = i * s1d1CtrlP_num + 1;
            asub[2] = i * s1d1CtrlP_num + 2;
            for(int i = 0; i < nzi; i++)
                A[row_idx*nv + asub[i]] = aval[i];
            row_idx++;
        }
    }      

    /*   End position  */
    {
        // position :
        for(int i = 0; i < 3; i++)
        {  // loop for x, y, z       
            int nzi = 1;
            unsigned int asub[nzi];
            double aval[nzi];
            asub[0] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num;
            aval[0] = 1.0 * lstScale;
            A[row_idx*nv + asub[0]] = aval[0];
            row_idx ++;
        }
        // velocity :
        for(int i = 0; i < 3; i++)
        { 
            int nzi = 2;
            unsigned int asub[nzi];
            double aval[nzi];
            asub[0] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num - 1;
            asub[1] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num;
            //  s(c1-c2) = 0    so,'s' can be 1
            aval[0] = - 1.0;
            aval[1] =   1.0;
            for(int i = 0; i < nzi; i++)
                A[row_idx*nv + asub[i]] = aval[i];
            row_idx++;
        }
        // acceleration : 
        for(int i = 0; i < 3; i++)
        { 
            int nzi = 3;
            unsigned int asub[nzi];
            double aval[nzi];
            asub[0] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num - 2;
            asub[1] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num - 1;
            asub[2] = ctrlP_num - 1 - (2 - i) * s1d1CtrlP_num;
            aval[0] =   1.0 / lstScale;
            aval[1] = - 2.0 / lstScale;
            aval[2] =   1.0 / lstScale;
            for(int i = 0; i < nzi; i++)
                A[row_idx*nv + asub[i]] = aval[i];
            row_idx++;
        }
    }

    /*   joint points  */
    if(segment_num > 1)
    {
        int sub_shift = 0;
        double val0, val1;
        for(int k = 0; k < (segment_num - 1); k ++ )
        {   
            double scale_k = corridor[k].t;
            double scale_n = corridor[k+1].t;
            // position :
            val0 = scale_k;
            val1 = scale_n;
            for(int i = 0; i < 3; i++)
            {  // loop for x, y, z
                int nzi = 2;
                unsigned int asub[nzi];
                double aval[nzi];

                // This segment's last control point
                aval[0] = 1.0 * val0;
                asub[0] = sub_shift + (i+1) * s1d1CtrlP_num - 1;

                // Next segment's first control point
                aval[1] = -1.0 * val1;
                asub[1] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num;
                for(int i = 0; i < nzi; i++)
                    A[row_idx*nv + asub[i]] = aval[i];
                row_idx++;
            }
            //velocity:
            for(int i = 0; i < 3; i++)
            {  
                int nzi = 4;
                unsigned int asub[nzi];
                double aval[nzi];
                
                // This segment's last velocity control point
                aval[0] = -1.0;
                aval[1] =  1.0;
                asub[0] = sub_shift + (i+1) * s1d1CtrlP_num - 2;    
                asub[1] = sub_shift + (i+1) * s1d1CtrlP_num - 1;   
                // Next segment's first velocity control point
                aval[2] =  1.0;
                aval[3] = -1.0;
                asub[2] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num;    
                asub[3] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num + 1;
                for(int i = 0; i < nzi; i++)
                    A[row_idx*nv + asub[i]] = aval[i];
                row_idx++;
            }
            // acceleration :
            val0 = 1.0 / scale_k;
            val1 = 1.0 / scale_n;
            for(int i = 0; i < 3; i++)
            {  
                int nzi = 6;
                unsigned int asub[nzi];
                double aval[nzi];
                
                // This segment's last velocity control point
                aval[0] =  1.0  * val0;
                aval[1] = -2.0  * val0;
                aval[2] =  1.0  * val0;
                asub[0] = sub_shift + (i+1) * s1d1CtrlP_num - 3;    
                asub[1] = sub_shift + (i+1) * s1d1CtrlP_num - 2;   
                asub[2] = sub_shift + (i+1) * s1d1CtrlP_num - 1;   
                // Next segment's first velocity control point
                aval[3] =  -1.0  * val1;
                aval[4] =   2.0  * val1;
                aval[5] =  -1.0  * val1;
                asub[3] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num;    
                asub[4] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num + 1;
                asub[5] = sub_shift + s1CtrlP_num + i * s1d1CtrlP_num + 2;
                for(int i = 0; i < nzi; i++)
                    A[row_idx*nv + asub[i]] = aval[i];
                row_idx++;
            }

            sub_shift += s1CtrlP_num;
        }
    }

    int min_order_l = floor(minimize_order);
    int min_order_u = ceil (minimize_order);

    //number of non-zero elements
    int NUMQNZ = 0;
    for(int i = 0; i < segment_num; i ++)
    {   
        int NUMQ_blk = (traj_order + 1);                       // default minimize the jerk and minimize_order = 3
        NUMQNZ      += 3 * NUMQ_blk * (NUMQ_blk + 1) / 2;
    }
    unsigned int qsubi[NUMQNZ], qsubj[NUMQNZ];
    double       qval[NUMQNZ];

    {    
        int sub_shift = 0;
        int idx = 0;
        for(int k = 0; k < segment_num; k ++)
        {
            double scale_k = corridor[k].t;
            for(int p = 0; p < 3; p ++ )
                for( int i = 0; i < s1d1CtrlP_num; i ++ )
                    for( int j = 0; j < s1d1CtrlP_num; j ++ )
                        if( i >= j )
                        {
                            qsubi[idx] = sub_shift + p * s1d1CtrlP_num + i;   
                            qsubj[idx] = sub_shift + p * s1d1CtrlP_num + j;  
                            //qval[idx]  = MQM(i, j) /(double)pow(scale_k, 3);
                            if(min_order_l == min_order_u)
                                qval[idx]  = MQM(i, j) /(double)pow(scale_k, 2 * min_order_u - 3);
                            else
                                qval[idx] = ( (minimize_order - min_order_l) / (double)pow(scale_k, 2 * min_order_u - 3)
                                            + (min_order_u - minimize_order) / (double)pow(scale_k, 2 * min_order_l - 3) ) * MQM(i, j);

//                            G[qsubi[idx]][qsubj[idx]] = qval[idx];
                            H[qsubi[idx]*nv + qsubj[idx]] = qval[idx];
//                            if(i != j)
                                H[qsubj[idx]*nv + qsubi[idx]] = qval[idx];
                            idx ++ ;
                        }

            sub_shift += s1CtrlP_num;
        }
    }

    QProblem solver(nv, nc);
    solver.setPrintLevel(PL_NONE);
//    solver.setPrintLevel(PL_MEDIUM);
    int_t nWSR = 600;
    real_t time = 1.0;
    real_t* cputime = &time;
    auto state = solver.init(H.get(), g.get(), A.get(), lb.get(), ub.get(), lbA.get(), ubA.get(), nWSR, cputime);

    /* Get and print solution of first QP. */
    real_t xopt[nv];
//    real_t yopt[nv+nc];
    if(state == SUCCESSFUL_RETURN){
        solver.getPrimalSolution(xopt);
//        solver.getDualSolution(yopt);
//        cout << "xOpt = [" ;
//        for(int i = 0; i < nv; i++){
//            cout << xopt[i] << ", ";
//        }
//        cout << "\b\b]" << endl;
//        cout << "\b\b]; yOpt = [";
//        for(int i = 0; i < nv+nc; i++){
//            cout << yopt[i] << ", ";
//        }
//        cout << "\b\b]; objVal = " << solver.getObjVal() << endl;
    } else {
        ROS_WARN("In solver, falied ");
        return -1;
    }

    VectorXd d_var(ctrlP_num);
    for(int i = 0; i < ctrlP_num; i++)
        d_var(i) = xopt[i];
    
    PolyCoeff = MatrixXd::Zero(segment_num, 3 *(traj_order + 1) );

    int var_shift = 0;
    for(int i = 0; i < segment_num; i++ )
    {
        for(int j = 0; j < 3 * n_poly; j++)
            PolyCoeff(i , j) = d_var(j + var_shift);

        var_shift += 3 * n_poly;
    }   

    return 1;
}
